package com.ideaaedi.component.compile;

import lombok.extern.slf4j.Slf4j;

import javax.tools.JavaFileObject;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 自定义类加载器，用于加载{@link CharSequenceJavaFileObject}实例
 *
 * @author JustryDeng
 * @since 2021/9/20 16:30:44
 */
@Slf4j
public class DynamicClassLoader extends ClassLoader {
    
    /**
     * key - CharSequenceJavaFileObject所代表类的全类名<br>
     * value - CharSequenceJavaFileObject实例
     */
    private final Map<String, CharSequenceJavaFileObject> javaFileObjectMap = new ConcurrentHashMap<>(16);
    
    /**
     * 假设现在我们要调用{@link DynamicClassLoader#findClass(String)}方法，defineClass两个类<br>
     * <code>public class  TestLogger1</code><br>
     * 和<br>
     * <code>public class TestLogger2 extends TestLogger1</code>时，<br>
     * <p>
     * 当先defineClass的是TestLogger2时，由于TestLogger2中用到了TestLogger1，那么defineClass方法内部就自动调用对TestLogger1的defineClass，<br>
     * 此时，调用findClass方法的调用者只收到了TestLogger2的Class实例，如果调用者不知道TestLogger2中有TestLogger1（即：不知道TestLogger1已经被<br>
     * defineClass了的话），调用者再主动对TestLogger1进行findClass时，就会因为重复defineClass，而报错：<br>
     * Caused by: java.lang.LinkageError: loader (instance of  com/ideaaedi/hot/clazz/compile/DynamicClassLoader): attempted  duplicate class definition for name: "com/test/TestLogger1"<br>
     * <p>
     * currDefinedClass参数就是来解决这个问题的
     */
    private final Map<String, Class<?>> currDefinedClass = new ConcurrentHashMap<>(8);
    
    public DynamicClassLoader(ClassLoader parentClassLoader) {
        super(parentClassLoader);
    }
    
    @Override
    protected Class<?> findClass(String classLongName) throws ClassNotFoundException {
        CharSequenceJavaFileObject javaFileObject = javaFileObjectMap.get(classLongName);
        if (javaFileObject != null) {
            byte[] byteCode = javaFileObject.getByteCode();
            Class<?> klass = defineClass(classLongName, byteCode, 0, byteCode.length);
            log.info("defineClass {} complement.", classLongName);
            currDefinedClass.put(classLongName, klass);
            return klass;
        }
        return super.findClass(classLongName);
    }
    
    @Override
    public InputStream getResourceAsStream(String name) {
        if (name.endsWith(JavaFileObject.Kind.CLASS.extension)) {
            String classLongNameWithoutSuffix = name.substring(0, name.length() - JavaFileObject.Kind.CLASS.extension.length()).replace('/', '.');
            CharSequenceJavaFileObject javaFileObject = javaFileObjectMap.get(classLongNameWithoutSuffix);
            if (javaFileObject != null && javaFileObject.getByteCode() != null) {
                return new ByteArrayInputStream(javaFileObject.getByteCode());
            }
        }
        return super.getResourceAsStream(name);
    }
    
    /**
     * 存
     */
    public void registerCompiledSource(String classLongName, CharSequenceJavaFileObject javaFileObject) {
        javaFileObjectMap.put(classLongName, javaFileObject);
    }
    
    /**
     * 罗列
     */
    public Collection<JavaFileObject> listCharSequenceJavaFileObject() {
        return Collections.unmodifiableCollection(javaFileObjectMap.values());
    }
    
    /**
     * 获取所有已加载的class<br/>
     * key - 全类名<br/>
     * value - class实例
     */
    public Map<String, Class<?>> getClasses() throws ClassNotFoundException {
        Map<String, Class<?>> classes = new HashMap<>(16);
        for (CharSequenceJavaFileObject byteCode : javaFileObjectMap.values()) {
            String classLongName = byteCode.getName();
            if (classes.containsKey(classLongName)) {
                log.info("classLongName already been defined.");
                continue;
            }
            Class<?> klass = findClass(classLongName);
            classes.putAll(currDefinedClass);
            classes.put(classLongName, klass);
            currDefinedClass.clear();
        }
        return classes;
    }
    
    /**
     * 获取所有已加载的class<br/>
     * key - 全类名<br/>
     * value - .class文件字节码
     */
    public Map<String, byte[]> getByteCodes() {
        Map<String, byte[]> result = new HashMap<>(javaFileObjectMap.size());
        for (Map.Entry<String, CharSequenceJavaFileObject> entry : javaFileObjectMap.entrySet()) {
            result.put(entry.getKey(), entry.getValue().getByteCode());
        }
        return result;
    }
}
